from SocketServer import ThreadingUDPServer, BaseRequestHandler
from socket import *
from struct import unpack

class DNSdHandler(BaseRequestHandler):
    def setup(self):
        self.timeout = 2
        self.upstream_server = ('8.8.8.8', 53)
        #self.upstream_server = ('208.67.222.222', 53)
        self.known_poisoned_address = (
            0x0807c62d,
            0x253d369e,
            0x2e52ae44,
            0x3b1803ad,
            0x4e10310f,
            0x5d2e0859,
            0x9f6a794b,
            0xcb620741,
            0xf3b9bb27,
            )

    def handle(self):
        req, socket = self.request
        socket.sendto(self.resolve(req), self.client_address)

    def resolve(self, req):
        s = socket(AF_INET, SOCK_DGRAM)
        s.connect(self.upstream_server)
        s.sendall(req)
        s.settimeout(self.timeout)
        while True:
            try:
                r = s.recv(0x10000)
                if not self.gfw_poisoned(r):
                    return r
            except timeout:
                return ''

    def resolve_TCP(self, req):
        # not implemented yet
        pass

    def gfw_poisoned(self, d):
        questions, answer_RRs, authority_RRs, Additional_RRs = unpack('>HHHH', d[4:12])
        Type, Class = unpack('>HH', d[-14:-10])
        if questions != 1 \
            or answer_RRs != 1 \
            or authority_RRs != 0 \
            or Additional_RRs != 0 \
            or Type != 1 \
            or Class != 1:
            return False
        address = unpack('>I', d[-4:])[0]
        """
        print questions, answer_RRs, authority_RRs, Additional_RRs, Type, Class, \
            len(d), \
            '0x%08x' % address, \
            '%d.%d.%d.%d' % unpack('>BBBB', d[-4:])
        """
        return address in self.known_poisoned_address

def main():
    DNSd = ThreadingUDPServer(('127.0.0.1', 53), DNSdHandler)
    DNSd.serve_forever()

if __name__ == '__main__':
    main()

